"""
AST序列化和反序列化

提供AST节点的序列化和反序列化功能，支持JSON、YAML等格式。
"""

import json
import yaml
from typing import Any, Dict, Type, Union
from abc import ABC, abstractmethod

from .base import Ast, AstNodeType
from .nodes import (
    Sequence, Parallel, Conditional, Loop,
    FfmpegSlice, ExtractMetadata, ScanDirectory,
    GeminiClassify, OutputDraft, OrganizeFiles, GenerateReport,
    FileOperation, DataTransform
)


class SerializationError(Exception):
    """序列化错误"""
    pass


class AstEncoder(ABC):
    """AST编码器抽象基类"""
    
    @abstractmethod
    def encode(self, ast: Ast) -> str:
        """编码AST为字符串"""
        pass


class AstDecoder(ABC):
    """AST解码器抽象基类"""
    
    @abstractmethod
    def decode(self, data: str) -> Ast:
        """从字符串解码AST"""
        pass


class JsonAstEncoder(AstEncoder):
    """JSON格式AST编码器"""
    
    def encode(self, ast: Ast) -> str:
        """编码AST为JSON字符串"""
        try:
            data = self._ast_to_dict(ast)
            return json.dumps(data, indent=2, ensure_ascii=False)
        except Exception as e:
            raise SerializationError(f"JSON编码失败: {e}")
    
    def _ast_to_dict(self, ast: Ast) -> Dict[str, Any]:
        """将AST转换为字典"""
        data = ast.to_dict()
        
        # 处理子节点
        if hasattr(ast, 'children') and ast.children:
            data['children'] = [self._ast_to_dict(child) for child in ast.children]
        
        return data


class JsonAstDecoder(AstDecoder):
    """JSON格式AST解码器"""
    
    # 节点类型映射
    NODE_TYPE_MAP = {
        AstNodeType.SEQUENCE: Sequence,
        AstNodeType.PARALLEL: Parallel,
        AstNodeType.CONDITIONAL: Conditional,
        AstNodeType.LOOP: Loop,
        AstNodeType.FFMPEG_SLICE: FfmpegSlice,
        AstNodeType.EXTRACT_METADATA: ExtractMetadata,
        AstNodeType.SCAN_DIRECTORY: ScanDirectory,
        AstNodeType.GEMINI_CLASSIFY: GeminiClassify,
        AstNodeType.OUTPUT_DRAFT: OutputDraft,
        AstNodeType.ORGANIZE_FILES: OrganizeFiles,
        AstNodeType.GENERATE_REPORT: GenerateReport,
        AstNodeType.FILE_OPERATION: FileOperation,
        AstNodeType.DATA_TRANSFORM: DataTransform,
    }
    
    def decode(self, data: str) -> Ast:
        """从JSON字符串解码AST"""
        try:
            json_data = json.loads(data)
            return self._dict_to_ast(json_data)
        except json.JSONDecodeError as e:
            raise SerializationError(f"JSON解码失败: {e}")
        except Exception as e:
            raise SerializationError(f"AST构建失败: {e}")
    
    def _dict_to_ast(self, data: Dict[str, Any]) -> Ast:
        """从字典构建AST"""
        # 获取节点类型
        node_type_str = data.get('node_type')
        if not node_type_str:
            raise SerializationError("缺少节点类型信息")
        
        try:
            node_type = AstNodeType(node_type_str)
        except ValueError:
            raise SerializationError(f"未知的节点类型: {node_type_str}")
        
        # 获取节点类
        node_class = self.NODE_TYPE_MAP.get(node_type)
        if not node_class:
            raise SerializationError(f"未注册的节点类型: {node_type}")
        
        # 构建节点
        metadata = data.get('metadata', {})
        
        # 处理特殊参数
        if node_type == AstNodeType.CONDITIONAL:
            # 条件节点需要特殊处理
            children_data = data.get('children', [])
            true_branch = self._dict_to_ast(children_data[0]) if len(children_data) > 0 else None
            false_branch = self._dict_to_ast(children_data[1]) if len(children_data) > 1 else None
            
            node = node_class(
                condition=metadata.get('condition', ''),
                true_branch=true_branch,
                false_branch=false_branch,
                **{k: v for k, v in metadata.items() if k != 'condition'}
            )
        else:
            # 普通节点
            node = node_class(**metadata)
        
        # 设置基础属性
        node.node_id = data.get('node_id', node.node_id)
        node.name = data.get('name', node.name)
        node.description = data.get('description', node.description)
        
        # 处理子节点（复合节点）
        if hasattr(node, 'children') and 'children' in data:
            children_data = data['children']
            if node_type != AstNodeType.CONDITIONAL:  # 条件节点已经处理过子节点
                node.children = [self._dict_to_ast(child_data) for child_data in children_data]
        
        return node


class YamlAstEncoder(AstEncoder):
    """YAML格式AST编码器"""
    
    def encode(self, ast: Ast) -> str:
        """编码AST为YAML字符串"""
        try:
            data = self._ast_to_dict(ast)
            return yaml.dump(data, default_flow_style=False, allow_unicode=True, indent=2)
        except Exception as e:
            raise SerializationError(f"YAML编码失败: {e}")
    
    def _ast_to_dict(self, ast: Ast) -> Dict[str, Any]:
        """将AST转换为字典"""
        data = ast.to_dict()
        
        # 处理子节点
        if hasattr(ast, 'children') and ast.children:
            data['children'] = [self._ast_to_dict(child) for child in ast.children]
        
        return data


class YamlAstDecoder(AstDecoder):
    """YAML格式AST解码器"""
    
    def __init__(self):
        self.json_decoder = JsonAstDecoder()
    
    def decode(self, data: str) -> Ast:
        """从YAML字符串解码AST"""
        try:
            yaml_data = yaml.safe_load(data)
            return self.json_decoder._dict_to_ast(yaml_data)
        except yaml.YAMLError as e:
            raise SerializationError(f"YAML解码失败: {e}")
        except Exception as e:
            raise SerializationError(f"AST构建失败: {e}")


class AstSerializationManager:
    """AST序列化管理器"""
    
    def __init__(self):
        self.encoders = {
            'json': JsonAstEncoder(),
            'yaml': YamlAstEncoder(),
        }
        self.decoders = {
            'json': JsonAstDecoder(),
            'yaml': YamlAstDecoder(),
        }
    
    def encode(self, ast: Ast, format: str = 'json') -> str:
        """编码AST"""
        encoder = self.encoders.get(format.lower())
        if not encoder:
            raise SerializationError(f"不支持的编码格式: {format}")
        
        return encoder.encode(ast)
    
    def decode(self, data: str, format: str = 'json') -> Ast:
        """解码AST"""
        decoder = self.decoders.get(format.lower())
        if not decoder:
            raise SerializationError(f"不支持的解码格式: {format}")
        
        return decoder.decode(data)
    
    def save_to_file(self, ast: Ast, file_path: str, format: str = None) -> None:
        """保存AST到文件"""
        if format is None:
            # 根据文件扩展名推断格式
            if file_path.endswith('.yaml') or file_path.endswith('.yml'):
                format = 'yaml'
            else:
                format = 'json'
        
        data = self.encode(ast, format)
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(data)
    
    def load_from_file(self, file_path: str, format: str = None) -> Ast:
        """从文件加载AST"""
        if format is None:
            # 根据文件扩展名推断格式
            if file_path.endswith('.yaml') or file_path.endswith('.yml'):
                format = 'yaml'
            else:
                format = 'json'
        
        with open(file_path, 'r', encoding='utf-8') as f:
            data = f.read()
        
        return self.decode(data, format)
    
    def register_encoder(self, format: str, encoder: AstEncoder) -> None:
        """注册编码器"""
        self.encoders[format.lower()] = encoder
    
    def register_decoder(self, format: str, decoder: AstDecoder) -> None:
        """注册解码器"""
        self.decoders[format.lower()] = decoder


# 全局序列化管理器实例
serialization_manager = AstSerializationManager()
